#!/usr/bin/python3
# coding=utf-8

import torch
import torch.nn.functional as F
import numpy as np
import os

def col2im_torch(col, output_size, kernel_size, dilation=1, stride=1, padding=0):
    """PyTorch实现的col2im函数
    
    参数:
        col: im2col展开后的矩阵，形状为 [N, C*kH*kW, L]
        output_size: 输出图像尺寸 (H, W)
        kernel_size: 卷积核大小，可以是int或tuple
        dilation: 空洞率，默认1
        stride: 步长，默认1
        padding: 填充大小，默认0
        
    返回:
        还原后的图像，形状为 [N, C, H, W]
    """
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    
    # 使用PyTorch的fold函数实现col2im，包含dilation参数
    x_reconstructed = F.fold(col, output_size, kernel_size, 
                            dilation=dilation, stride=stride, padding=padding)
    return x_reconstructed

def gen_golden_data_col2im():
    """生成col2im操作的测试数据"""
    # 配置参数
    dtype = torch.float32
    output_shape = [9, 10, 23, 1971]  # [N, C, H, W] - 这是col2im的输出形状
    
    # 卷积参数
    kernel_size = (3, 2)   # 卷积核大小
    dilation = 3           # 空洞率
    stride = 5             # 步长
    padding = 4            # 填充大小
    
    # 计算col矩阵的形状（考虑dilation）
    N, C, H, W = output_shape
    kH, kW = kernel_size
    dH, dW = dilation, dilation
    
    # 考虑dilation的有效核尺寸
    effective_kH = (kH - 1) * dH + 1
    effective_kW = (kW - 1) * dW + 1
    
    out_H = (H + 2 * padding - effective_kH) // stride + 1
    out_W = (W + 2 * padding - effective_kW) // stride + 1
    L = out_H * out_W
    output_channels = C * kH * kW
    col_shape = [N, output_channels, L]
    
    # 生成col输入数据 - 使用有规律的数据便于验证
    torch.manual_seed(42)
    col = torch.rand(col_shape, dtype=dtype)
    
    # 创建输入输出目录
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # 计算col2im结果（包含dilation）
    output_size = (H, W)
    golden = col2im_torch(col, output_size, kernel_size, dilation, stride, padding)
    
    # 转换为numpy用于保存
    col_np = col.numpy()
    golden_np = golden.numpy()
    
    # 保存二进制文件
    col_np.tofile("./input/input_col.bin")
    golden_np.tofile("./output/golden.bin")

if __name__ == "__main__":
    gen_golden_data_col2im()
